import torch
import torch.nn as nn


# 1. 单通道池化
def test01():
    # 定义输入输数据 【1,3,3 】
    inputs = torch.tensor([[[0, 1, 2], [3, 4, 5], [6, 7, 8]]]).float()
    # 修改stride，padding观察效果
    # 1. 最大池化
    polling = nn.MaxPool2d(kernel_size=2, stride=1, padding=0)
    output = polling(inputs)
    print("最大池化：\n", output)
    # 2. 平均池化
    polling = nn.AvgPool2d(kernel_size=2, stride=1, padding=0)
    output = polling(inputs)
    print("平均池化：\n", output)


# 2. 多通道池化
def test02():
    # 定义输入输数据 【3,3,3 】
    inputs = torch.tensor([[[0, 1, 2], [3, 4, 5], [6, 7, 8]],
                           [[10, 20, 30], [40, 50, 60], [70, 80, 90]],
                           [[11, 22, 33], [44, 55, 66], [77, 88, 99]]]).float()
    # 最大池化
    polling = nn.MaxPool2d(kernel_size=2, stride=1, padding=0)
    output = polling(inputs)
    print("多通道池化：\n", output)


if __name__ == '__main__':
    # test01()
    test02()
